#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import time

sys.path.append('.')

"""
    手工过滤不符合要求的图片的可视化工具
"""


h_image_count = 10
v_image_count = 10
img_width = 96
img_height = 96

batch_size = 32


ix, iy = -1, -1
draw_image = None
image_bboxes = []
select_image_bbox_idxes = set()


def get_xy_bbox_idx(x, y):
    global image_bboxes
    bbox_idx = None
    for i, bbox in enumerate(image_bboxes):
        # print(bbox, x, y)
        if (bbox[0] < x < bbox[2]) and (bbox[1] < y < bbox[3]):
            bbox_idx = i
            break
    return bbox_idx


def IOU(bbox, bboxes):
    bboxes2 = np.array(bboxes, dtype=np.int32)
    bboxes1 = np.array([bbox], dtype=np.int32)
    # bboxes1 = np.array([bbox], dtype=np.int32).repeat(bboxes2.shape[0], axis=0)
    lt = np.maximum(bboxes1[:, :2], bboxes2[:, :2])
    rb = np.minimum(bboxes1[:, 2:], bboxes2[:, 2:])
    wh = (rb - lt)
    wh = np.maximum(wh, 0)
    overlap = wh[:, 0] * wh[:, 1]
    area1 = (bboxes1[:, 2] - bboxes1[:, 0]) * (bboxes1[:, 3] - bboxes1[:, 1])
    area2 = (bboxes2[:, 2] - bboxes2[:, 0]) * (bboxes2[:, 3] - bboxes2[:, 1])
    ious = overlap / (area1 + area2 - overlap)
    return ious


def get_xy_bbox_idxes(x1, y1, x2, y2):
    global image_bboxes
    if x1 < x2:
        sx = x1
        ex = x2
    else:
        sx = x2
        ex = x1
    if y1 < y2:
        sy = y1
        ey = y2
    else:
        sy = y2
        ey = y1
    iou = IOU([sx, sy, ex, ey], image_bboxes)
    bbox_idxes = np.nonzero(iou > 0)[0].tolist()
    return bbox_idxes


def draw_rect(image, start_xy, end_xy):
    if (start_xy is None) or (end_xy is None):
        return
    x1, y1 = start_xy
    x2, y2 = end_xy
    if x1 < x2:
        sx = x1
        ex = x2
    else:
        sx = x2
        ex = x1
    if y1 < y2:
        sy = y1
        ey = y2
    else:
        sy = y2
        ey = y1
    cv2.rectangle(image, (sx, sy), (ex, ey), (0, 255, 0), thickness=2)


start_bbox_init = False
start_modify_bbox = False
start_mouse_xy = None
end_mouse_xy = None
start_bbox_init2 = False
start_modify_bbox2 = False
start_mouse_xy2 = None
end_mouse_xy2 = None


def mouse_event(event, x, y, flags, param):
    global draw_image, image_bboxes, select_image_bbox_idxes, start_bbox_init, start_modify_bbox, start_mouse_xy, \
        end_mouse_xy, start_bbox_init2, start_modify_bbox2, start_mouse_xy2, end_mouse_xy2

    if draw_image is not None:
        if event == cv2.EVENT_MOUSEMOVE and flags == (cv2.EVENT_FLAG_SHIFTKEY + cv2.EVENT_FLAG_LBUTTON):
            if start_modify_bbox:
                end_mouse_xy = [x, y]
                bbox_idxes = get_xy_bbox_idxes(start_mouse_xy[0], start_mouse_xy[1], x, y)
                for bbox_idx in bbox_idxes:
                    if bbox_idx not in select_image_bbox_idxes:
                        select_image_bbox_idxes.add(bbox_idx)
            return

        if event == cv2.EVENT_LBUTTONDOWN and flags == (cv2.EVENT_FLAG_SHIFTKEY + cv2.EVENT_FLAG_LBUTTON):
            if not start_bbox_init:
                start_bbox_init = True
                start_modify_bbox = True
                start_mouse_xy = [x, y]
            return

        start_bbox_init = False
        start_modify_bbox = False
        start_mouse_xy = None
        end_mouse_xy = None

        if event == cv2.EVENT_MOUSEMOVE and flags == (cv2.EVENT_FLAG_CTRLKEY + cv2.EVENT_FLAG_LBUTTON):
            if start_modify_bbox2:
                end_mouse_xy2 = [x, y]
                bbox_idxes = get_xy_bbox_idxes(start_mouse_xy2[0], start_mouse_xy2[1], x, y)
                for bbox_idx in bbox_idxes:
                    if bbox_idx in select_image_bbox_idxes:
                        select_image_bbox_idxes.remove(bbox_idx)
            return

        if event == cv2.EVENT_LBUTTONDOWN and flags == (cv2.EVENT_FLAG_CTRLKEY + cv2.EVENT_FLAG_LBUTTON):
            if not start_bbox_init2:
                start_bbox_init2 = True
                start_modify_bbox2 = True
                start_mouse_xy2 = [x, y]
            return

        start_bbox_init2 = False
        start_modify_bbox2 = False
        start_mouse_xy2 = None
        end_mouse_xy2 = None

        if event == cv2.EVENT_LBUTTONDOWN:
            # print(x, y)
            bbox_idx = get_xy_bbox_idx(x, y)
            if bbox_idx is not None:
                if bbox_idx not in select_image_bbox_idxes:
                    select_image_bbox_idxes.add(bbox_idx)
                else:
                    select_image_bbox_idxes.remove(bbox_idx)


def read_file_list(file_list, use_select=True):
    file_count = 0
    files = []
    selects = []
    file_selects = []
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            lines = line.strip().split()

            file_name = lines[0]
            if use_select:
                select = int(lines[1])
            else:
                select = 0
            file_selects.append([file_name, select])
    file_selects.sort(key=lambda x: x[0])
    for file_name, select in file_selects:
        files.append(file_name)
        selects.append(select)

        print('Read File Count: {}'.format(file_count))

    return files, selects


def save_process(select_file_list_path, select_file_list, cur_img_idx_file_path, cur_img_idx):
    if len(select_file_list) > 0:
        with open(select_file_list_path, 'w') as file:
            select_file_names, is_selects = select_file_list
            for select_file_name, is_select in zip(select_file_names, is_selects):
                select_info = '{} {}\n'.format(select_file_name, is_select)
                # print(select_info)
                file.write(select_info)
        os.system('chmod a+wr {}'.format(select_file_list_path))

    with open(cur_img_idx_file_path, 'w') as file:
        process_info = '{}\n'.format(cur_img_idx)
        file.write(process_info)
    os.system('chmod a+wr {}'.format(cur_img_idx_file_path))


def filter_select_face(filter_file_list_path, select_file_list):
    filter_file_list_file = open(filter_file_list_path, 'w')
    face_count = 0
    face_image_count = 0
    files, selects = select_file_list
    has_face_image = False
    for file_name, is_select in zip(files, selects):
        if 0 != is_select:
            continue
        out_info = '{}\n'.format(file_name)
        # print(out_info)
        filter_file_list_file.write(out_info)
        face_image_count += 1
        has_face_image = True
    if has_face_image:
        face_count += 1
    filter_file_list_file.close()
    os.system('chmod a+wr {}'.format(filter_file_list_path))
    print('Face Count: {}, Face Image Count: {}'.format(face_count, face_image_count))


def main():
    global draw_image, image_bboxes, select_image_bbox_idxes, start_modify_bbox, start_mouse_xy, end_mouse_xy, \
        start_modify_bbox2, start_mouse_xy2, end_mouse_xy2
    parser = argparse.ArgumentParser(description="Select Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data1/data/VideoCropFace", help="path to image file", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/VideoCropFace/face_image_q95-100.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/VideoCropFace/mask_select", help="path to mot reid image info", type=str
    )
    parser.add_argument("--use_temp", action='store_true', help="use temp save, if crash can use it")
    parser.add_argument('--save_time', type=int, default=3*60, help='save time, second ')
    parser.add_argument("--filter", action='store_true', help="filter select image")
    args = parser.parse_args()

    save_time = args.save_time
    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    display_image_count = h_image_count * v_image_count

    base_name = os.path.splitext(os.path.basename(args.file_list))[0]

    select_file_list_path = os.path.join(output_dir, base_name + '_select_images.txt')
    if not os.path.exists(select_file_list_path):
        select_file_list = read_file_list(args.file_list, False)
    else:
        select_file_list = read_file_list(select_file_list_path)

    if args.filter:
        filter_file_list_path = os.path.join(output_dir, base_name + '_filter.txt')
        filter_select_face(filter_file_list_path, select_file_list)
        return

    cv2.namedWindow('image')
    cv2.setMouseCallback('image', mouse_event)

    cur_img_idx = 0
    cur_img_idx_file_path = os.path.join(output_dir, base_name + '_cur_img_idx.txt')
    if os.path.exists(cur_img_idx_file_path):
        with open(cur_img_idx_file_path, 'r') as file:
            for line in file.readlines():
                cur_img_idx = int(line.strip())
                break

    if args.use_temp:
        tmp_output_dir = os.path.join('/dev/shm', 'SelectImage')
        if not os.path.exists(tmp_output_dir):
            os.makedirs(tmp_output_dir)
            os.system('chmod a+wr {}'.format(tmp_output_dir))

        tmp_select_file_list_path = os.path.join(tmp_output_dir, base_name + '_select_images.txt')
        tmp_cur_id_idx_file_path = os.path.join(tmp_output_dir, base_name + '_cur_img_idx.txt')

    select_change = False
    start_save_time = time.time()

    file_list, select_list = select_file_list

    last_idx = len(file_list) - 1
    start_idx = 0
    files = []

    try:
        idx = cur_img_idx
        while True:
            file_name = file_list[idx]
            files.append(file_name)
            file_count = len(files)
            if file_count == 1:
                start_idx = idx
            if file_count < display_image_count and (idx < last_idx):
                idx += 1
                continue
            idx += 1
            imgs = []
            for file_name in files:
                file_path = os.path.join(image_dir, file_name)
                img = cv2.imread(file_path)

                # img = cv2.resize(img, (img_width, img_height))

                w, h = img.shape[1], img.shape[0]
                wscale = img_width / w
                hscale = img_height / h
                scale = min(wscale, hscale)
                w = int(round(w * scale))
                h = int(round(h * scale))
                image = cv2.resize(img, (w, h))
                img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
                x = (img_width-w)//2
                y = (img_height-h)//2
                img[y:y+h, x:x+w, :] = image

                imgs.append(img)

            image_count = len(imgs)
            cur_img_idx = start_idx + image_count
            images = []
            x_off = 0
            y_off = 0
            image_bboxes.clear()
            select_image_bbox_idxes.clear()
            for select_idx in range(0, image_count):
                if 0 != select_list[start_idx+select_idx]:
                    select_image_bbox_idxes.add(select_idx)
            for j, img in enumerate(imgs):
                if j % h_image_count == 0:
                    x_off = 0
                    if j != 0:
                        y_off += img_height
                    h_images = [img]
                else:
                    x_off += img_width
                    h_images.append(img)
                image_bboxes.append([x_off, y_off, (x_off + img_width), (y_off + img_height)])
                if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                    images.append(np.concatenate(h_images, axis=1))

            if len(images) > 1:
                w0 = images[0].shape[1]
                h0 = images[0].shape[0]
                w1 = images[-1].shape[1]
                h1 = images[-1].shape[0]
                if (w0 != w1) or (h0 != h1):
                    # image = np.empty_like(images[0])
                    image = np.zeros_like(images[0])
                    image[0:h1, 0:w1, :] = images[-1]
                    images[-1] = image
            draw_image = np.vstack(images)

            is_exit = False
            while True:
                image = draw_image.copy()
                if start_modify_bbox:
                    draw_rect(image, start_mouse_xy, end_mouse_xy)
                if start_modify_bbox2:
                    draw_rect(image, start_mouse_xy2, end_mouse_xy2)
                for select_idx in select_image_bbox_idxes:
                    bbox = image_bboxes[select_idx]
                    cv2.circle(image, (bbox[0]+14, bbox[1]+14), 8, (0, 0, 255), -1)

                cv2.putText(image, '%d/%d' % (idx, last_idx + 1),
                            (8, 24), cv2.FONT_HERSHEY_PLAIN, 2.0, (0, 255, 0), thickness=2)
                cv2.imshow('image', image)
                key = cv2.waitKey(50)
                if (key == 27) or (key == ord('q')) or (key == ord('Q')):
                    is_exit = True
                    break
                elif key == ord('a') or key == ord('A'):  # 上一页
                    select_change = True
                    if start_idx >= display_image_count:
                        idx = start_idx - display_image_count
                    break
                elif key == ord(' ') or key == ord('d') or key == ord('D'):  # 下一页
                    select_change = True
                    break
                elif key == ord('e') or key == ord('E'):  # 当前页全选
                    select_image_bbox_idxes.update(range(0, image_count))
                    continue
                elif key == ord('w') or key == ord('W'):  # 当前页全不选
                    select_image_bbox_idxes.clear()
                    continue
                elif key == ord('s') or key == ord('S'):  # Save
                    if args.use_temp:
                        save_process(tmp_select_file_list_path, select_file_list, tmp_cur_id_idx_file_path, cur_img_idx)
                    save_process(select_file_list_path, select_file_list, cur_img_idx_file_path, cur_img_idx)
                    start_save_time = time.time()
                    select_change = False
                    continue
                if args.use_temp:
                    # 定时保存
                    if time.time() - start_save_time >= save_time:
                        if select_change:
                            save_process(tmp_select_file_list_path, select_file_list, tmp_cur_id_idx_file_path, cur_img_idx)
                            select_change = False
                        start_save_time = time.time()
            # if len(select_image_bbox_idxes) > 0:
            for bbox_idx in range(0, image_count):
                select_idx = start_idx + bbox_idx
                if bbox_idx in select_image_bbox_idxes:
                    select_list[select_idx] = 1
                else:
                    select_list[select_idx] = 0
            files.clear()
            if is_exit:
                save_process(select_file_list_path, select_file_list, cur_img_idx_file_path, cur_img_idx)
                return
    except Exception as e:
        print(e)
    save_process(select_file_list_path, select_file_list, cur_img_idx_file_path, cur_img_idx)
    print('finish!')


if __name__ == '__main__':
    main()
